package com.kawins.supers.message.socket.manager;

import org.springframework.web.socket.WebSocketSession;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentSkipListMap;

/**
 * @author rankai
 *         createTime 2017-10-2017/10/30 17:03
 */
public final class WebSocketBank {

    private static final Map<Serializable, Map<Serializable, WebSocketSession>> BANK_MAP = new ConcurrentSkipListMap<>();

    private WebSocketBank() {

    }

    public static void put(Serializable key, Serializable sessionId, WebSocketSession webSocketSession) {
        Map<Serializable, WebSocketSession> hashMap = BANK_MAP.get(key);
        if (hashMap == null) {
            hashMap = new HashMap<>(1);
            BANK_MAP.put(key, hashMap);
        }
        hashMap.put(sessionId, webSocketSession);
    }

    public static List<WebSocketSession> get(Serializable key) {
        Map<Serializable, WebSocketSession> hashMap = BANK_MAP.get(key);
        if (hashMap != null) {
            List<WebSocketSession> list = new ArrayList(hashMap.size());
            for (WebSocketSession webSocketSession : hashMap.values()) {
                list.add(webSocketSession);
            }
            return list;
        }
        return null;
    }

    public static WebSocketSession get(Serializable key, Serializable sessionId) {
        Map<Serializable, WebSocketSession> hashMap = BANK_MAP.get(key);
        if (hashMap != null) {
            return hashMap.get(sessionId);
        }
        return null;
    }

    public static List<WebSocketSession> getAll() {
        List<WebSocketSession> list = new ArrayList(0);
        for (Map<Serializable, WebSocketSession> map : BANK_MAP.values()) {
            for (WebSocketSession session : map.values()) {
                list.add(session);
            }
        }
        return null;
    }

    public static void remove(Serializable key) {
        BANK_MAP.remove(key);
    }

    public static void remove(Serializable key, Serializable sessionId) {
        Map<Serializable, WebSocketSession> hashMap = BANK_MAP.get(key);
        if (hashMap != null) {
            if (hashMap.size() == 1) {
                remove(key);
            } else {
                hashMap.remove(sessionId);
            }
        }
    }
}
